# -*- coding: utf-8 -*-

"""antlr4 为每个语法规则提供了访问方法 visitXXXX(ctx)，罗列全面，方便重写
antlr语法分析器的精髓是为每个节点都提供了同名的节点对象，该节点对象还可以用子节点名称直接访问子节点对象

！！！始终牢记自顶向下下降算法：一个语法分析程序，生成自己的AST节点；上级算法调用下级算法，下级算法返回的节点作为上级算法节点的子节点。
直接通过 evaluator.visitProgram(ProgramContext::=AST) [顶级算法(整个AST)] 实现对脚本生成的AST的遍历计算
"""

from antlr4.tree.Tree import TerminalNodeImpl
from HelloSyntaxParser import HelloSyntaxParser
from HelloSyntaxVisitor import HelloSyntaxVisitor

class HelloASTEvaluator(HelloSyntaxVisitor):
    def __init__(self) -> None:
        super().__init__()
        self.variables = {}

    # Visit a parse tree produced by HelloSyntaxParser#program.
    def visitProgram(self, ctx:HelloSyntaxParser.ProgramContext):
        results = []
        for stmt_ctx in ctx.children:
            results.append(self.visitExpressionStmt(stmt_ctx))
        return results

    # Visit a parse tree produced by HelloSyntaxParser#expressionStmt.
    def visitExpressionStmt(self, ctx:HelloSyntaxParser.ExpressionStmtContext):
        if child:= ctx.forLoopStmt():
            return self.visitForLoopStmt(child)
        
        if child:= ctx.ifStmt():
            return self.visitIfStmt(child)

        if child:= ctx.intDeclaration():
            return self.visitIntDeclaration(child)

        if child:= ctx.assignmentExpression():
            return self.visitAssignmentExpression(child)

        if child:= ctx.boolExpression():
            return self.visitBoolExpression(child)


    # Visit a parse tree produced by HelloSyntaxParser#forLoopStmt.
    def visitForLoopStmt(self, ctx:HelloSyntaxParser.ForLoopStmtContext):
        self.visitIntDeclaration(ctx.intDeclaration())
        condition = self.visitBoolExpression(ctx.boolExpression())
        while(condition):
            for stmt_ctx in ctx.expressionStmt():
                self.visitExpressionStmt(stmt_ctx)

            self.visitAssignmentExpression(ctx.assignmentExpression())
            condition = self.visitBoolExpression(ctx.boolExpression())


    # Visit a parse tree produced by HelloSyntaxParser#ifStmt.
    def visitIfStmt(self, ctx:HelloSyntaxParser.IfStmtContext):
        condition = self.visitBoolExpression(ctx.boolExpression())
        if condition:
            if isinstance(ctx.success, list):
                for expressionStmt in ctx.success:
                    self.visitExpressionStmt(expressionStmt)
            else:
                self.visitExpressionStmt(ctx.success)
        else:
            if isinstance(ctx.faild, list):
                for expressionStmt in ctx.faild:
                    self.visitExpressionStmt(expressionStmt)
            else:
                self.visitExpressionStmt(ctx.faild)


    # Visit a parse tree produced by HelloSyntaxParser#intDeclaration.
    def visitIntDeclaration(self, ctx:HelloSyntaxParser.IntDeclarationContext):
        name, val = ctx.Identifier().getText(), None

        if child:= ctx.boolExpression():
            val = self.visitBoolExpression(child)
        self.variables[name] = val

        return {name: val}


    # Visit a parse tree produced by HelloSyntaxParser#assignmentExpression.
    def visitAssignmentExpression(self, ctx:HelloSyntaxParser.AssignmentExpressionContext):
        name, val = ctx.Identifier().getText(), self.visitBoolExpression(ctx.boolExpression())
        self.variables[name] = val

        return {name: val}


    # Visit a parse tree produced by HelloSyntaxParser#boolExpression.
    def visitBoolExpression(self, ctx:HelloSyntaxParser.BoolExpressionContext):
        comp_rs_l = self.visitComparision(ctx.comparision(0))
        
        if ctx.And() or ctx.Or():
            bool_type = None
            children = ctx.children
            for child in children[1:]:
                if isinstance(child, TerminalNodeImpl):
                    bool_type = child.getSymbol().type
                else:
                    comp_r = self.visitComparision(child)
                    
                    if bool_type is HelloSyntaxParser.And:
                        comp_rs_l = comp_rs_l and comp_r
                    else:
                        comp_rs_l = comp_rs_l or comp_r
            
        return comp_rs_l


    # Visit a parse tree produced by HelloSyntaxParser#comparision.
    def visitComparision(self, ctx:HelloSyntaxParser.ComparisionContext):
        addi_l = self.visitAdditiveEx(ctx.additiveEx(0))
        if len(ctx.children) > 1:
            addi_r = self.visitAdditiveEx(ctx.additiveEx(1))
            if ctx.GT():
                return addi_l > addi_r
            elif ctx.GE():
                return addi_l >= addi_r
            elif ctx.EQ():
                return addi_l == addi_r
            elif ctx.LT():
                return addi_l < addi_r
            elif ctx.LE():
                return addi_l <= addi_r
        else:
            return addi_l


    # Visit a parse tree produced by HelloSyntaxParser#additiveEx.
    def visitAdditiveEx(self, ctx:HelloSyntaxParser.AdditiveExContext):
        multi_rs_l = self.visitMultiplicativeEX(ctx.multiplicativeEX(0))
        
        if ctx.Plus() or ctx.Minus():
            addi_type = None
            children = ctx.children
            for child in children[1:]:
                if isinstance(child, TerminalNodeImpl):
                    addi_type = child.getSymbol().type
                else:
                    multi_rs_r = self.visitMultiplicativeEX(child)
                    if addi_type is HelloSyntaxParser.Plus:
                        multi_rs_l += multi_rs_r
                    else:
                        multi_rs_l -= multi_rs_r

        return multi_rs_l


    # Visit a parse tree produced by HelloSyntaxParser#multiplicativeEX.
    def visitMultiplicativeEX(self, ctx:HelloSyntaxParser.MultiplicativeEXContext):
        prim_l = self.visitPrimary(ctx.primary(0))
        
        if ctx.Star() or ctx.Slash():
            multi_type = None
            children = ctx.children
            for child in children[1:]:
                if isinstance(child, TerminalNodeImpl):
                    multi_type = child.getSymbol().type
                else:
                    prim_r = self.visitPrimary(child)
                    if multi_type is HelloSyntaxParser.Star:
                        prim_l *= prim_r
                    else:
                        prim_l /= prim_r

        return prim_l


    # Visit a parse tree produced by HelloSyntaxParser#primary.
    def visitPrimary(self, ctx:HelloSyntaxParser.PrimaryContext):
        prim_type = ctx.children[0].getSymbol().type
        if prim_type is HelloSyntaxParser.Identifier:
            name = ctx.Identifier().getText()
            return self.variables[name]
        elif prim_type is HelloSyntaxParser.IntLiteral:
            return int(ctx.IntLiteral().getText())
        elif prim_type is HelloSyntaxParser.StringLiteral:
            return ctx.StringLiteral().getText()
        else:
            return self.visitBoolExpression(ctx.boolExpression())

